﻿
open System
open System.IO
open System.Text
open System.Collections.Generic

open fsgen.types
open fsgen.opcodes

type Immediate =
    | FuncIdx
    | LocalIdx
    | GlobalIdx
    | LabelIdx
    | ResultType
    | I32
    | I64
    | F32
    | F64
    | U32
    | U8
    | MemArg
    | CallIndirect
    | BrTable
    | Nothing

let add_immediate (d: Dictionary<string,Immediate>) name im =
    d.Add(name, im)

let verify_immediates (d: Dictionary<string,Immediate>) =
    let h = 
        let h = HashSet<string>()
        for op in opcode_infos do
            h.Add(op.name) |> ignore
        h
    for s in d.Keys do
        if not (h.Contains(s)) then
            printfn "INVALID opcode in immediate lookup table: %s" s

let build_immediate_lookup () =
    let d = Dictionary<string,Immediate>()
    add_immediate d "I32Const" I32
    add_immediate d "I64Const" I64
    add_immediate d "F32Const" F32
    add_immediate d "F64Const" F64
    add_immediate d "Call" FuncIdx
    add_immediate d "Br" LabelIdx
    add_immediate d "BrIf" LabelIdx
    add_immediate d "LocalGet" LocalIdx
    add_immediate d "LocalSet" LocalIdx
    add_immediate d "LocalTee" LocalIdx
    add_immediate d "GlobalGet" GlobalIdx
    add_immediate d "GlobalSet" GlobalIdx
    add_immediate d "Block" ResultType
    add_immediate d "Loop" ResultType
    add_immediate d "If" ResultType
    add_immediate d "I32Load" MemArg
    add_immediate d "I64Load" MemArg
    add_immediate d "F32Load" MemArg
    add_immediate d "F64Load" MemArg
    add_immediate d "I32Load8S" MemArg
    add_immediate d "I32Load8U" MemArg
    add_immediate d "I32Load16S" MemArg
    add_immediate d "I32Load16U" MemArg
    add_immediate d "I64Load8S" MemArg
    add_immediate d "I64Load8U" MemArg
    add_immediate d "I64Load16S" MemArg
    add_immediate d "I64Load16U" MemArg
    add_immediate d "I64Load32S" MemArg
    add_immediate d "I64Load32U" MemArg
    add_immediate d "I32Store" MemArg
    add_immediate d "I64Store" MemArg
    add_immediate d "F32Store" MemArg
    add_immediate d "F64Store" MemArg
    add_immediate d "I32Store8" MemArg
    add_immediate d "I32Store16" MemArg
    add_immediate d "I64Store8" MemArg
    add_immediate d "I64Store16" MemArg
    add_immediate d "I64Store32" MemArg
    add_immediate d "MemorySize" U8 // TODO what is this arg?
    add_immediate d "MemoryGrow" U8 // TODO what is this arg?
    add_immediate d "CallIndirect" CallIndirect
    add_immediate d "BrTable" BrTable

    verify_immediates d

    d

let get_immediate (d: Dictionary<string,Immediate>) name =
    if (d.ContainsKey(name)) then
        d.[name]
    else
        Nothing

let get_prefixes () =
    let h = HashSet<int>()
    for op in opcode_infos do
        match op.prefix with
        | Some p -> h.Add(p) |> ignore
        | None -> ()
    h

let write_type_instruction path (immediates: Dictionary<string,Immediate>) =
    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.def_instr" |> pr
    "    open wasm.def_basic" |> pr
    "    type Instruction =" |> pr
    for op in opcode_infos do
        match get_immediate immediates op.name with
        | FuncIdx -> sprintf "        | %s of FuncIdx"  op.name |> pr
        | LocalIdx -> sprintf "        | %s of LocalIdx"  op.name |> pr
        | GlobalIdx -> sprintf "        | %s of GlobalIdx"  op.name |> pr
        | LabelIdx -> sprintf "        | %s of LabelIdx"  op.name |> pr
        | ResultType -> sprintf "        | %s of ValType option"  op.name |> pr
        | I32 -> sprintf "        | %s of int32"  op.name |> pr
        | I64 -> sprintf "        | %s of int64"  op.name |> pr
        | F32 -> sprintf "        | %s of float32"  op.name |> pr
        | F64 -> sprintf "        | %s of double"  op.name |> pr
        | U32 -> sprintf "        | %s of uint32"  op.name |> pr
        | U8  -> sprintf "        | %s of byte"  op.name |> pr
        | MemArg  -> sprintf "        | %s of MemArg"  op.name |> pr
        | CallIndirect  -> sprintf "        | %s of CallIndirectArg"  op.name |> pr
        | BrTable  -> sprintf "        | %s of BrTableArg"  op.name |> pr
        | Nothing -> sprintf "        | %s"  op.name |> pr
    "\n" |> pr
    let txt = sb.ToString()
    File.WriteAllText(path, txt)
    
let write_function_read_instruction path (immediates: Dictionary<string,Immediate>) =
    let prefixes = get_prefixes()

    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.read_instr" |> pr
    "    open wasm.read_basic" |> pr
    "    open wasm.def_basic" |> pr
    "    open wasm.def_instr" |> pr
    "    open wasm.read_args" |> pr

    pr "    let read_instruction br ="
    pr "        let b1 = read_byte br"
    pr "        match b1 with"
    for n in prefixes do
        sprintf "        | 0x%02xuy ->" n |> pr
        sprintf "            let b2 = read_byte br" |> pr
        sprintf "            match b2 with" |> pr
        for op in opcode_infos do
            match op.prefix with
            | Some x -> 
                if x = n then
                    sprintf "            | 0x%02xuy -> %s" op.code op.name |> pr
            | None -> ()
        sprintf "            | _      -> failwith \"unknown opcode\"" |> pr

    for op in opcode_infos do
        match op.prefix with
        | Some _ -> ()
        | None -> 
            match get_immediate immediates op.name with
            | FuncIdx -> sprintf "        | 0x%02xuy -> %s (read_var_u32 br |> FuncIdx)" op.code op.name |> pr
            | LocalIdx -> sprintf "        | 0x%02xuy -> %s (read_var_u32 br |> LocalIdx)" op.code op.name |> pr
            | GlobalIdx -> sprintf "        | 0x%02xuy -> %s (read_var_u32 br |> GlobalIdx)" op.code op.name |> pr
            | LabelIdx -> sprintf "        | 0x%02xuy -> %s (read_var_u32 br |> LabelIdx)" op.code op.name |> pr
            | ResultType -> sprintf "        | 0x%02xuy -> %s (match read_byte br with | 0x40uy -> None | x -> x |> make_valtype |> Some)" op.code op.name |> pr
            | I32 -> sprintf "        | 0x%02xuy -> %s (read_var_i32 br)" op.code op.name |> pr
            | I64 -> sprintf "        | 0x%02xuy -> %s (read_var_i64 br)" op.code op.name |> pr
            | F32 -> sprintf "        | 0x%02xuy -> %s (read_f32 br)" op.code op.name |> pr
            | F64 -> sprintf "        | 0x%02xuy -> %s (read_f64 br)" op.code op.name |> pr
            | U8 -> sprintf "        | 0x%02xuy -> %s (read_byte br)" op.code op.name |> pr
            | U32 -> sprintf "        | 0x%02xuy -> %s (read_var_u32 br)" op.code op.name |> pr
            | MemArg -> sprintf "        | 0x%02xuy -> %s (read_memarg br)" op.code op.name |> pr
            | CallIndirect -> sprintf "        | 0x%02xuy -> %s (read_callindirect br)" op.code op.name |> pr
            | BrTable -> sprintf "        | 0x%02xuy -> %s (read_brtable br)" op.code op.name |> pr
            | Nothing -> sprintf "        | 0x%02xuy -> %s" op.code op.name |> pr
    sprintf "        | _      -> failwith \"unknown opcode\"" |> pr
    sprintf "" |> pr

    let txt = sb.ToString()
    File.WriteAllText(path, txt)
    
let write_function_write_instruction path (immediates: Dictionary<string,Immediate>) =
    let prefixes = get_prefixes()

    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.write_instr" |> pr
    "    open wasm.def_basic" |> pr
    "    open wasm.def_instr" |> pr
    "    open wasm.write_basic" |> pr
    "    open wasm.write_args" |> pr

    pr "    let write_instruction (w: System.IO.BinaryWriter) op ="
    pr "        match op with"
    for op in opcode_infos do
        match op.prefix with
        | Some _ -> ()
        | None -> 
            let imm = get_immediate immediates op.name

            match imm with
            | Nothing -> sprintf "        | %s ->" op.name |> pr
            | _ -> sprintf "        | %s i ->" op.name |> pr

            sprintf "            write_byte w 0x%02xuy" op.code |> pr

            match imm with
            | FuncIdx -> sprintf "            let (FuncIdx i) = i in write_var_u32 w i" |> pr
            | LocalIdx -> sprintf "            let (LocalIdx i) = i in write_var_u32 w i" |> pr
            | GlobalIdx -> sprintf "            let (GlobalIdx i) = i in write_var_u32 w i" |> pr
            | LabelIdx -> sprintf "            let (LabelIdx i) = i in write_var_u32 w i" |> pr
            | ResultType -> sprintf "            match i with | Some vt -> write_byte w (encode_valtype vt) | None -> write_byte w 0x40uy" |> pr
            | I32 -> sprintf "            write_var_i32 w i" |> pr
            | I64 -> sprintf "            write_var_i64 w i" |> pr
            | F32 -> sprintf "            write_f32 w i" |> pr
            | F64 -> sprintf "            write_f64 w i" |> pr
            | U8 -> sprintf "            write_byte w i" |> pr
            | U32 -> sprintf "            write_var_u32 w i" |> pr
            | MemArg -> sprintf "            write_memarg w i" |> pr
            | CallIndirect -> sprintf "            write_callindirect w i" |> pr
            | BrTable -> sprintf "            write_brtable w i" |> pr
            | Nothing -> ()
    sprintf "" |> pr

    let txt = sb.ToString()
    File.WriteAllText(path, txt)
    
(*
let write_function_stringify_instruction path (immediates: Dictionary<string,Immediate>) =
    let prefixes = get_prefixes()

    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.stringify_instr" |> pr
    "    open wasm.def_basic" |> pr
    "    open wasm.def_instr" |> pr
    "    open wasm.stringify_args" |> pr

    pr "    let stringify_instruction funcs op ="
    pr "        match op with"
    for op in opcode_infos do
        match op.prefix with
        | Some _ -> ()
        | None -> 
            let imm = get_immediate immediates op.name

            match imm with
            | Nothing -> sprintf "        | %s -> \"%s\"" op.name op.text |> pr
            | I32 -> sprintf "        | %s i -> sprintf \"%s %%d\" i" op.name op.text |> pr
            | I64 -> sprintf "        | %s i -> sprintf \"%s %%d\" i" op.name op.text |> pr
            | F32 -> sprintf "        | %s i -> sprintf \"%s %%f\" i" op.name op.text |> pr
            | F64 -> sprintf "        | %s i -> sprintf \"%s %%f\" i" op.name op.text |> pr
            | U8 -> sprintf "        | %s i -> sprintf \"%s %%u\" i" op.name op.text |> pr
            | U32 -> sprintf "        | %s i -> sprintf \"%s %%u\" i" op.name op.text |> pr
            | FuncIdx -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_funcidx i)" op.name op.text |> pr
            | LocalIdx -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_localidx i)" op.name op.text |> pr
            | GlobalIdx -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_globalidx i)" op.name op.text |> pr
            | LabelIdx -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_labelidx i)" op.name op.text |> pr
            | ResultType -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_resulttype i)" op.name op.text |> pr
            | MemArg -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_memarg i)" op.name op.text |> pr
            | CallIndirect -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_callindirect i)" op.name op.text |> pr
            | BrTable -> sprintf "        | %s i -> sprintf \"%s %%s\" (funcs.stringify_brtable i)" op.name op.text |> pr

    sprintf "" |> pr

    let txt = sb.ToString()
    File.WriteAllText(path, txt)
*)
    
let write_function_instruction_name path (immediates: Dictionary<string,Immediate>) =
    let prefixes = get_prefixes()

    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.instr_name" |> pr
    "    open wasm.def_basic" |> pr
    "    open wasm.def_instr" |> pr

    pr "    let get_instruction_name op ="
    pr "        match op with"
    for op in opcode_infos do
        match op.prefix with
        | Some _ -> ()
        | None -> 
            let imm = get_immediate immediates op.name

            match imm with
            | Nothing -> sprintf "        | %s -> \"%s\"" op.name op.text |> pr
            | _ -> sprintf "        | %s _ -> \"%s\"" op.name op.text |> pr

    sprintf "" |> pr

    let txt = sb.ToString()
    File.WriteAllText(path, txt)
    
let write_function_instruction_stack_info path (immediates: Dictionary<string,Immediate>) =
    let prefixes = get_prefixes()

    let sb = StringBuilder()
    let pr (s: string) =
        sb.Append(s + "\n") |> ignore
    "// this file is automatically generated" |> pr
    "module wasm.instr_stack" |> pr
    "    open wasm.def_basic" |> pr
    "    open wasm.def_instr" |> pr

    pr "    let get_instruction_stack_info op ="
    pr "        match op with"

    let str_result t =
        match t with
        | Some q -> sprintf "Some %A" q
        | None -> "None"

    let get_stack_info op =
        match (op.type1, op.type2) with
        | (Some a1, Some a2) -> sprintf "TwoArgs { rtype = %s; arg1 = %A; arg2 = %A; }" (str_result op.rtype) a1 a2
        | (Some a1, None) -> sprintf "OneArg { rtype = %s; arg = %A; }" (str_result op.rtype) a1 
        | (None, Some a2) -> failwith "illegal"
        | (None, None) -> sprintf "NoArgs (%s)" (str_result op.rtype)

    for op in opcode_infos do
        match op.prefix with
        | Some _ -> ()
        | None -> 
            match op.name with
            | "Call" -> sprintf "        | %s i -> SpecialCaseCall i" op.name |> pr
            | "CallIndirect" -> sprintf "        | %s i -> SpecialCaseCallIndirect i" op.name |> pr
            | "LocalGet" -> sprintf "        | %s i -> SpecialCaseLocalGet i" op.name |> pr
            | "LocalSet" -> sprintf "        | %s i -> SpecialCaseLocalSet i" op.name |> pr
            | "LocalTee" -> sprintf "        | %s i -> SpecialCaseLocalTee i" op.name |> pr
            | "GlobalGet" -> sprintf "        | %s i -> SpecialCaseGlobalGet i" op.name |> pr
            | "GlobalSet" -> sprintf "        | %s i -> SpecialCaseGlobalSet i" op.name |> pr
            | "Drop" -> sprintf "        | %s -> SpecialCaseDrop" op.name |> pr
            | "Select" -> sprintf "        | %s -> SpecialCaseSelect" op.name |> pr
            | "Block" -> sprintf "        | %s t -> SpecialCaseBlock t" op.name |> pr
            | "If" -> sprintf "        | %s t -> SpecialCaseIf t" op.name |> pr
            | "Else" -> sprintf "        | %s -> SpecialCaseElse" op.name |> pr
            | "Loop" -> sprintf "        | %s t -> SpecialCaseLoop t" op.name |> pr
            | "End" -> sprintf "        | %s -> SpecialCaseEnd" op.name |> pr
            | "Br" -> sprintf "        | %s _ -> SpecialCaseBr" op.name |> pr
            | "BrTable" -> sprintf "        | %s _ -> SpecialCaseBrTable" op.name |> pr
            | "Return" -> sprintf "        | %s -> SpecialCaseReturn" op.name |> pr
            | "Unreachable" -> sprintf "        | %s -> SpecialCaseUnreachable" op.name |> pr
            
            // er, the following rows in opcode.def are incorrect, so fix them here
            | "BrIf" -> sprintf "        | %s _ -> OneArg { rtype = None; arg = I32; }" op.name |> pr
            | "F32Sqrt" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Sqrt" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Abs" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Abs" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Neg" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Neg" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Ceil" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Ceil" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Floor" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Floor" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Trunc" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Trunc" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr
            | "F32Nearest" -> sprintf "        | %s -> OneArg { rtype = Some F32; arg = F32; }" op.name |> pr
            | "F64Nearest" -> sprintf "        | %s -> OneArg { rtype = Some F64; arg = F64; }" op.name |> pr

            | _ ->
                let imm = get_immediate immediates op.name
                match imm with
                | Nothing -> sprintf "        | %s -> %s" op.name (get_stack_info op) |> pr
                | _ -> sprintf "        | %s _ -> %s" op.name (get_stack_info op) |> pr

    sprintf "" |> pr

    let txt = sb.ToString()
    File.WriteAllText(path, txt)
[<EntryPoint>]
let main argv =
    let dir_top =
        let cwd = Directory.GetCurrentDirectory()
        Path.GetFullPath(Path.Combine(cwd, ".."))
        // TODO or maybe walk upward until we find the right directory

    let dir_wasm = Path.Combine(dir_top, "wasm")
    let immediates = build_immediate_lookup ()
    write_type_instruction (Path.Combine(dir_wasm, "Instruction.fs")) immediates
    write_function_read_instruction (Path.Combine(dir_wasm, "ReadInstruction.fs")) immediates
    write_function_write_instruction (Path.Combine(dir_wasm, "WriteInstruction.fs")) immediates
    //write_function_stringify_instruction (Path.Combine(dir_wasm, "StringifyInstruction.fs")) immediates
    write_function_instruction_name (Path.Combine(dir_wasm, "InstructionName.fs")) immediates
    write_function_instruction_stack_info (Path.Combine(dir_wasm, "InstructionStackInfo.fs")) immediates

    0 // return an integer exit code

